#!/usr/bin/python3

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from numpy import in1d, concatenate
from numpy import array as np_array
from numpy.random import randint
from numpy import arange as np_arange

from torch import sqrt, Tensor, LongTensor, from_numpy, stack
from torch import cat as torch_cat

from torch.utils.data import Dataset

class TrainDataset(Dataset):
    def __init__(self, triples, nentity, nrelation, negative_sample_size, mode):
        self.len = len(triples)
        self.triples = triples
        self.triple_set = set(triples)
        self.nentity = nentity
        self.nrelation = nrelation
        self.negative_sample_size = negative_sample_size
        self.mode = mode
        self.count = self.count_frequency(triples)
        self.true_head, self.true_tail = self.get_true_head_and_tail(self.triples)

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        positive_sample = self.triples[idx]

        head, relation, tail = positive_sample

        subsampling_weight = self.count[(head, relation)] + self.count[(tail, -relation - 1)]
        subsampling_weight = sqrt(1 / Tensor([subsampling_weight]))

        negative_sample_list = []
        negative_sample_size = 0

        while negative_sample_size < self.negative_sample_size:
            negative_sample = randint(self.nentity, size=self.negative_sample_size * 2)
            if self.mode == 'head-batch':
                mask = in1d(
                    negative_sample,
                    self.true_head[(relation, tail)],
                    assume_unique=True,
                    invert=True
                )
            elif self.mode == 'tail-batch':
                mask = in1d(
                    negative_sample,
                    self.true_tail[(head, relation)],
                    assume_unique=True,
                    invert=True
                )
            else:
                raise ValueError('Training batch mode %s not supported' % self.mode)
            negative_sample = negative_sample[mask]
            negative_sample_list.append(negative_sample)
            negative_sample_size += negative_sample.size

        negative_sample = concatenate(negative_sample_list)[:self.negative_sample_size]

        negative_sample = from_numpy(negative_sample).long()

        positive_sample = LongTensor(positive_sample)

        return positive_sample, negative_sample, subsampling_weight, self.mode, False

    @staticmethod
    def collate_fn(data):
        positive_sample = stack([_[0] for _ in data], dim=0)
        negative_sample = stack([_[1] for _ in data], dim=0)
        subsample_weight = torch_cat([_[2] for _ in data], dim=0)
        mode = data[0][3]
        return positive_sample, negative_sample, subsample_weight, mode

    @staticmethod
    def count_frequency(triples, start=4):
        '''
        Get frequency of a partial triple like (head, relation) or (relation, tail)
        The frequency will be used for subsampling like word2vec
        '''
        count = {}
        for head, relation, tail in triples:
            if (head, relation) not in count:
                count[(head, relation)] = start
            else:
                count[(head, relation)] += 1

            if (tail, -relation - 1) not in count:
                count[(tail, -relation - 1)] = start
            else:
                count[(tail, -relation - 1)] += 1
        return count

    @staticmethod
    def get_true_head_and_tail(triples):
        '''
        Build a dictionary of true triples that will
        be used to filter these true triples for negative sampling
        '''

        true_head = {}
        true_tail = {}

        for head, relation, tail in triples:
            if (head, relation) not in true_tail:
                true_tail[(head, relation)] = []
            true_tail[(head, relation)].append(tail)
            if (relation, tail) not in true_head:
                true_head[(relation, tail)] = []
            true_head[(relation, tail)].append(head)

        for relation, tail in true_head:
            true_head[(relation, tail)] = np_array(list(set(true_head[(relation, tail)])))
        for head, relation in true_tail:
            true_tail[(head, relation)] = np_array(list(set(true_tail[(head, relation)])))

        return true_head, true_tail


class TestDataset(Dataset):
    def __init__(self, triples, all_true_triples, nentity, nrelation, mode):
        self.len = len(triples)
        self.triple_set = set(all_true_triples)
        self.triples = triples
        self.nentity = nentity
        self.nrelation = nrelation
        self.mode = mode

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        head, relation, tail = self.triples[idx]
        if self.mode == 'head-batch':
            tmp = [(0, rand_head) if (rand_head, relation, tail) not in self.triple_set
                   else (-1, head) for rand_head in range(self.nentity)]
            try:
                tmp[head] = (0, head)
            except:
                print(head)
        elif self.mode == 'tail-batch':
            tmp = [(0, rand_tail) if (head, relation, rand_tail) not in self.triple_set
                   else (-1, tail) for rand_tail in range(self.nentity)]
            tmp[tail] = (0, tail)
        else:
            raise ValueError('negative batch mode %s not supported' % self.mode)

        tmp = LongTensor(tmp)
        filter_bias = tmp[:, 0].float()
        negative_sample = tmp[:, 1]

        positive_sample = LongTensor((head, relation, tail))

        return positive_sample, negative_sample, filter_bias, self.mode

    @staticmethod
    def collate_fn(data):
        positive_sample = stack([_[0] for _ in data], dim=0).long()
        negative_sample = stack([_[1] for _ in data], dim=0).long()
        filter_bias = stack([_[2] for _ in data], dim=0)
        mode = data[0][3]
        return positive_sample, negative_sample, filter_bias, mode


class BidirectionalOneShotIterator(object):
    def __init__(self, dataloader_head, dataloader_tail):
        self.iterator_head = self.one_shot_iterator(dataloader_head)
        self.iterator_tail = self.one_shot_iterator(dataloader_tail)
        self.step = 0

    def __next__(self):
        self.step += 1
        if self.step % 2 == 0:
            data = next(self.iterator_head)
        else:
            data = next(self.iterator_tail)
        return data

    @staticmethod
    def one_shot_iterator(dataloader):
        '''
        Transform a PyTorch Dataloader into python iterator
        '''
        while True:
            for data in dataloader:
                yield data


from multiprocessing import Pool

class UpdateBidirectionalIterator(object):
    def __init__(self, dataloader_head, dataloader_tail):
        self.iterator_head = dataloader_head
        self.iterator_tail = dataloader_tail
        self.step = 0

    def __next__(self):
        self.step += 1
        if self.step % 2 == 0:
            return next(self.iterator_head)
        return next(self.iterator_tail)


class UpdateDataIterator:

    def __init__(self, dataset, target_triples, n_workers, collect_fn):
        self.dataset = dataset
        self.n_workers = n_workers
        self.collect_fn = collect_fn
        self.target_triples = target_triples

    def __next__(self):
        if len(self.target_triples) >= 100 and self.n_workers > 1:
            results = []
            pool = Pool(processes=self.n_workers)
            n_triples_per_worker = len(self.target_triples) // self.n_workers
            for i in range(self.n_workers - 1):
                results.append(pool.apply_async(self.sub_batch_sample,
                                               args=(self.target_triples[i*n_triples_per_worker: (i+1)*n_triples_per_worker],)))
            results.append(pool.apply_async(self.sub_batch_sample,
                                           args=(self.target_triples[(self.n_workers - 1) * n_triples_per_worker:],)))
            pool.close()
            pool.join()

            samples = []
            for result in results:
                samples.append(result.get())
        else:
            samples = [self.sub_batch_sample(self.target_triples)]

        positive_samples, negative_samples, subsample_weights = [], [], []
        for sample in samples:
            positive_samples.append(sample[0])
            negative_samples.append(sample[1])
            subsample_weights.append(sample[2])
        mode = samples[0][3]
        positive_samples = torch_cat(positive_samples).long()
        negative_samples = torch_cat(negative_samples).long()
        subsample_weights = torch_cat(subsample_weights)
        return positive_samples, negative_samples, subsample_weights, mode

    def sub_batch_sample(self, triples):
        data = []
        for triple in triples:
            data.append(self.dataset.getItem(triple))
        data = self.collect_fn(data)
        return data


class UpdateDataset:
    def __init__(self, triples, n_entity, negative_sample_size, mode):
        self.triples = triples          # np.array
        self.n_entity = n_entity
        self.negative_sample_size = negative_sample_size
        self.mode = mode
        self.count = self.count_frequency(triples)
        self.true_head, self.true_tail = self.get_true_head_and_tail(triples)

    def getItem(self, positive_sample):

        head, relation, tail = positive_sample

        # subsampling_weight = self.count.setdefault((head, relation), 0) + self.count.setdefault((tail, -relation - 1), 0)
        subsampling_weight = Tensor([1.])

        negative_sample_list = []
        negative_sample_size = 0

        while negative_sample_size < self.negative_sample_size:
            negative_sample = randint(self.n_entity, size=self.negative_sample_size * 2,)
            if self.mode == 'head-batch':
                try:
                    mask = in1d(
                        negative_sample,
                        self.true_head[(relation, tail)],
                        assume_unique=True,
                        invert=True
                    )
                except:
                    mask = np_arange(negative_sample.shape[0])
            elif self.mode == 'tail-batch':
                try:
                    mask = in1d(
                        negative_sample,
                        self.true_tail[(head, relation)],
                        assume_unique=True,
                        invert=True
                    )
                except:
                    mask = np_arange(negative_sample.shape[0])
            else:
                raise ValueError('Training batch mode %s not supported' % self.mode)
            negative_sample = negative_sample[mask]
            negative_sample_list.append(negative_sample)
            negative_sample_size += negative_sample.size

        negative_sample = concatenate(negative_sample_list)[:self.negative_sample_size]
        negative_sample = from_numpy(negative_sample)
        positive_sample = LongTensor(positive_sample)
        return positive_sample, negative_sample, subsampling_weight, self.mode, False

    @staticmethod
    def collate_fn(data):
        positive_sample = stack([_[0] for _ in data], dim=0)
        negative_sample = stack([_[1] for _ in data], dim=0)
        subsample_weight = torch_cat([_[2] for _ in data], dim=0)
        mode = data[0][3]
        return positive_sample, negative_sample, subsample_weight, mode

    @staticmethod
    def count_frequency(triples, start=4):
        '''
        Get frequency of a partial triple like (head, relation) or (relation, tail)
        The frequency will be used for subsampling like word2vec
        '''
        count = {}
        for head, relation, tail in triples:
            if (head, relation) not in count:
                count[(head, relation)] = start
            else:
                count[(head, relation)] += 1

            if (tail, -relation - 1) not in count:
                count[(tail, -relation - 1)] = start
            else:
                count[(tail, -relation - 1)] += 1
        return count

    @staticmethod
    def get_true_head_and_tail(triples):
        '''
        Build a dictionary of true triples that will
        be used to filter these true triples for negative sampling
        '''

        true_head = {}
        true_tail = {}

        for head, relation, tail in triples:
            if (head, relation) not in true_tail:
                true_tail[(head, relation)] = []
            true_tail[(head, relation)].append(tail)
            if (relation, tail) not in true_head:
                true_head[(relation, tail)] = []
            true_head[(relation, tail)].append(head)

        for relation, tail in true_head:
            true_head[(relation, tail)] = np_array(list(set(true_head[(relation, tail)])))
        for head, relation in true_tail:
            true_tail[(head, relation)] = np_array(list(set(true_tail[(head, relation)])))

        return true_head, true_tail


if __name__ == "__main__":
    triples = [(i, i+1, i+2) for i in range(100)]
    target_triples = [(3, 5, 7), (5, 7, 9), (7, 9, 11)]
    n_entity = 102
    n_workers = 2
    negative_sample_size = 4

    dataset_head = UpdateDataset(triples, n_entity, negative_sample_size, "head-batch")
    dataset_tail = UpdateDataset(triples, n_entity, negative_sample_size, "tail-batch")

    data_iterator_head = UpdateDataIterator(dataset_head, target_triples, n_workers, UpdateDataset.collate_fn)
    data_iterator_tail = UpdateDataIterator(dataset_tail, target_triples, n_workers, UpdateDataset.collate_fn)

    dataloader = UpdateBidirectionalIterator(data_iterator_head, data_iterator_tail)

    for _ in range(10):
        print(next(dataloader))